import collections
import datetime
import itertools
import random
import time
import traceback

from absl import app, flags

import numpy

import coin.proto.coin_order_enums_pb2 as coin_order
# TODO: consolidate better way..
from coin.strategy.kr_hard_arb import hard_arb_general
from coin.strategy.kr_hard_arb.arb_info import create_arb_info_general


def get_coin_diff(info, currency):
  total_coin = sum(hard_arb_general.get_currency_balance_map(info, currency).values())
  expected_total_coin = info.config["qty"]
  return expected_total_coin - total_coin


def get_fill_vwap_price(orders):
  if len(orders) == 0:
    return None
  else:
    return numpy.average([order.price for order in orders], weights=[order.qty for order in orders])


def filter_fills(fills, side, min_time):
  fills_out = []
  for fill in fills:
    if fill.fill_timestamp > min_time and fill.side == side:
      fills_out.append(fill)
  return fills_out


def get_all_orders_map(info, currency):
  orders = collections.OrderedDict()
  for exch in info.exchs:
    orders[exch] = []
    list_orders = exch.order.query_list_orders(exch.order.get_product(currency))
    for order in list_orders.msg.each_order:
      orders[exch].append(order)
  return orders


def get_all_fills_map(info, currency):
  fills = collections.OrderedDict()
  for exch in info.exchs:
    fills[exch] = []
    list_fills = exch.order.client.query_fills(exch.order.get_product(currency))
    for fill in list_fills.msg.each_fill:
      fills[exch].append(fill)
  return fills


def get_price_det_map(info, currency):
  price_det_notional = info.config["price_det_notional"]
  price_det_map = collections.OrderedDict()
  for exch in info.exchs:
    price_det = exch.feed.get_price_by_notional(price_det_notional)
    if price_det is None:
      print("Cannot get price by notional, %s" % exch.exchange)
      return None
    price_det_map[exch.exchange] = price_det
  return price_det_map


def get_min_ask_max_bid(price_det_map):
  asks = []
  bids = []
  for price_det in price_det_map.values():
    asks.append(price_det.ask_price)
    bids.append(price_det.bid_price)
  return min(asks), max(bids)


def get_posting_info(info, side, self_ask, self_bid, min_ask, max_bid):
  # ex) maintain ask on [min_ask + 40bp, min_ask + 60bp]
  # and cancel all ask < self_ask + 60bp (as mm posting margin)
  # ex) maintain bid on [max_bid - 60bp, max_bid - 40bp]
  # and cancel all bid > self_bid - 60bp (as mm posting margin)
  tick_size = info.config["tick_size"]
  if side == coin_order.BUY_ORDER:
    price_max_all = min_ask * (1 - info.config["cancel_inside"])
    price_max_self = self_ask * (1 - info.config["self_margin"])
    posting_info = {
        'price_min': min_ask * (1 - info.config["cancel_outside"]),
    }
    posting_info['price_max'] = min(price_max_all, price_max_self)

    price_post = min_ask * (1 - info.config["rate_post"])
    # post cheap, floor.
    price_post = numpy.floor(price_post / tick_size) * tick_size
    print(posting_info, "buy", price_post, "self bid", self_bid)
    if (posting_info['price_min'] < price_post < posting_info['price_max']
        and price_post >= self_bid):
      posting_info['price_post'] = price_post
    else:
      posting_info['price_post'] = None
  elif side == coin_order.SELL_ORDER:
    price_min_all = max_bid * (1 + info.config["cancel_inside"])
    price_min_self = self_bid * (1 + info.config["self_margin"])
    posting_info = {
        'price_max': max_bid * (1 + info.config["cancel_outside"]),
    }
    posting_info['price_min'] = max(price_min_all, price_min_self)

    price_post = max_bid * (1 + info.config["rate_post"])
    # post expensive, ceil.
    price_post = numpy.ceil(price_post / tick_size) * tick_size
    print(posting_info, "sell", price_post, "self ask", self_ask)
    if (posting_info['price_min'] < price_post < posting_info['price_max']
        and price_post <= self_ask):
      posting_info['price_post'] = price_post
    else:
      posting_info['price_post'] = None
  else:
    raise ValueError(side)
  return posting_info


def should_cancel_order(info, order, self_ask, self_bid, min_ask, max_bid):
  posting_info = get_posting_info(info, order.side, self_ask, self_bid, min_ask, max_bid)
  return not (posting_info['price_min'] <= order.price <= posting_info['price_max'])


def manage_orders(info, currency, exch_orders_map, price_det_map):
  min_ask, max_bid = get_min_ask_max_bid(price_det_map)
  cancel_at_least_one = False
  for exch, orders in exch_orders_map.items():
    price_dict = {
        'self_ask': price_det_map[exch.exchange].ask_price,
        'self_bid': price_det_map[exch.exchange].bid_price,
        'min_ask': min_ask,
        'max_bid': max_bid,
    }
    print("%-10s" % exch.exchange, price_dict)
    product = exch.order.get_product(currency)
    for side in [coin_order.BUY_ORDER, coin_order.SELL_ORDER]:
      orders_sided = [order for order in orders if order.side == side]
      if len(orders_sided) == 0:
        continue
      order_prices = [order.price for order in orders_sided]
      order_max = orders_sided[order_prices.index(max(order_prices))]
      order_min = orders_sided[order_prices.index(min(order_prices))]
      # actually targets to cancel all orders once get out of boundary.
      # aims to work one order for each side only per exchange.
      # cannot maintain many orders anyway, due to the throughput limits..
      if (should_cancel_order(info, order_min, **price_dict)
          or should_cancel_order(info, order_max, **price_dict)):
        cancel_at_least_one = True
        exch.order.query_cancel_product(product)
        break
  return cancel_at_least_one


def send_order(info, side, currency, exch, exch_krw, exch_ccy, order_price, order_qty):
  print("%-10s Price %s" % (exch.exchange, order_price))
  order_qty = float(numpy.floor(order_qty * 1e4) * 1e-4)
  if side == coin_order.BUY_ORDER:
    print('Buy Side, Qty %s' % order_qty)
    if hard_arb_general.has_enough_bal(exch_krw, order_qty * order_price):
      exch.order.query_submit(exch.order.get_product(currency), 'bid', order_price, order_qty)
      return True
    else:
      return False
      print("Not enough Balance")
  elif side == coin_order.SELL_ORDER:
    print('Sell Side, Qty %s' % order_qty)
    if hard_arb_general.has_enough_bal(exch_ccy, order_qty):
      exch.order.query_submit(exch.order.get_product(currency), 'ask', order_price, order_qty)
      return True
    else:
      return False
      print("Not enough Balance")


def get_buy_sell_qty(info, currency):
  config = info.config
  lot_size_krw = config["lot_size_krw"]

  total_ccy = sum(hard_arb_general.get_currency_balance_map(info, currency).values())
  ref_price = hard_arb_general.get_ref_price(info, currency)

  lot_size = (lot_size_krw / float(ref_price)) * (1 + random.random() * 0.2)

  buy_qty = hard_arb_general.truncate(min(lot_size, max(config["max_qty"] - total_ccy, 0)))
  sell_qty = hard_arb_general.truncate(min(lot_size, total_ccy))
  return (buy_qty, sell_qty)


def post_orders(info, currency, exch_orders_map, price_det_map):
  min_ask, max_bid = get_min_ask_max_bid(price_det_map)
  for exch, orders in exch_orders_map.items():
    if not info.config["passive_post"][exch.exchange]:
      continue
    price_dict = {
        'self_ask': price_det_map[exch.exchange].ask_price,
        'self_bid': price_det_map[exch.exchange].bid_price,
        'min_ask': min_ask,
        'max_bid': max_bid
    }
    exch_krw = hard_arb_general.get_currency_balance(exch.order, 'KRW')
    exch_ccy = hard_arb_general.get_currency_balance(exch.order, currency)
    for side in [coin_order.BUY_ORDER, coin_order.SELL_ORDER]:
      posting_info = get_posting_info(info, side, **price_dict)
      orders_sided = [order for order in orders if order.side == side]
      # do not post if order exists
      if len(orders_sided) > 0:
        continue

      buy_qty, sell_qty = get_buy_sell_qty(info, currency)

      print("Post")

      if posting_info['price_post'] is None:
        continue

      return send_order(info,
                        side,
                        currency,
                        exch,
                        exch_krw,
                        exch_ccy,
                        posting_info['price_post'],
                        buy_qty if side == coin_order.BUY_ORDER else sell_qty)


def trade_arb(info, currency):
  coin_diff = get_coin_diff(info, currency)

  price_det_map = get_price_det_map(info, currency)
  if price_det_map is None:
    return

  # TODO: change to fetched time
  order_min_time = time.time() * 1e9 - info.config["arb_time_window"]

  exch_orders_map = get_all_orders_map(info, currency)
  exch_fills_map = get_all_fills_map(info, currency)

  # Manage First.
  if manage_orders(info, currency, exch_orders_map, price_det_map):
    # if managed, must wait next turn (throughput limit)
    return

  # Underfill
  if coin_diff < -info.config["qty_threshold"]:
    closed_exch = None
    filtered_fills = filter_fills(itertools.chain(*exch_fills_map.values()),
                                  coin_order.SELL_ORDER,
                                  order_min_time)
    ref_vwap_price = get_fill_vwap_price(filtered_fills)
    if ref_vwap_price is None:
      print("no ref price available")
      print(filtered_fills)
    else:
      ref_price = ref_vwap_price * (1 - info.config['rate_close'])
      for exch in info.exchs:
        # buy cheaper then last sell
        if price_det_map[exch.exchange].ask_price < ref_price:
          exch_krw = hard_arb_general.get_currency_balance(exch.order, 'KRW')
          exch_ccy = hard_arb_general.get_currency_balance(exch.order, currency)
          if send_order(info,
                        coin_order.BUY_ORDER,
                        currency,
                        exch,
                        exch_krw,
                        exch_ccy,
                        price_det_map[exch.exchange].ask_price,
                        -coin_diff):
            closed_exch = exch.exchange
            break
      print("previous sell order vwap %s" % (ref_vwap_price))
      print("target %s" % (ref_price))
      if closed_exch is not None:
        print("closed with %s" % closed_exch)
      else:
        print("closure failure")
  # Overfill
  elif coin_diff > info.config["qty_threshold"]:
    filtered_fills = filter_fills(itertools.chain(*exch_fills_map.values()),
                                  coin_order.BUY_ORDER,
                                  order_min_time)
    ref_vwap_price = get_fill_vwap_price(filtered_fills)
    if ref_vwap_price is None:
      print("no ref price available")
      print(filtered_fills)
    else:
      ref_price = ref_vwap_price * (1 + info.config['rate_close'])
      closed_exch = None
      for exch in info.exchs:
        if price_det_map[exch.exchange].bid_price > ref_price:
          exch_krw = hard_arb_general.get_currency_balance(exch.order, 'KRW')
          exch_ccy = hard_arb_general.get_currency_balance(exch.order, currency)
          if send_order(info,
                        coin_order.SELL_ORDER,
                        currency,
                        exch,
                        exch_krw,
                        exch_ccy,
                        price_det_map[exch.exchange].bid_price,
                        coin_diff):
            closed_exch = exch.exchange
            break
      print("previous buy order vwap %s" % (ref_vwap_price))
      print("target %s" % (ref_price))
      if closed_exch is not None:
        print("closed with %s" % closed_exch)
      else:
        print("closure failure")
  post_orders(info, currency, exch_orders_map, price_det_map)


def trade(info, currency):
  try:
    print()
    hard_arb_general.execute_thread_runners(
        [hard_arb_general.ThreadRunner(exch.feed.query_book) for exch in info.exchs]
        + [hard_arb_general.ThreadRunner(exch.order.query_balance) for exch in info.exchs],
        "Collecting %s Data" % currency)
    hard_arb_general.print_status(info, currency)
    trade_arb(info, currency)
  except KeyboardInterrupt:
    raise
  except Exception:
    print("A problem detected. Try again")
    traceback.print_exc()


def main(_):
  currency = flags.FLAGS.currency

  info = create_arb_info_general(currency,
                                 flags.FLAGS.exchanges.split(","),
                                 flags.FLAGS.key_files.split(","),
                                 flags.FLAGS.config_filename)

  assert info.config["arb_time_window"] > 0
  assert info.config["qty_threshold"] >= 0
  assert info.config["rate_close"] >= 0
  assert info.config["cancel_outside"] > info.config["cancel_inside"]
  assert info.config["self_margin"] >= 0
  assert "passive_post" in info.config
  for exch in info.exchs:
    assert exch.exchange in info.config["passive_post"]

  while (True):
    time.sleep(2)
    print()
    print(datetime.datetime.now())
    print("Checking...")
    trade(info, currency)
  return 0


if __name__ == '__main__':
  app.run(main)
